from typing import List
from task_allocation  import task_allocation_kdtree,reallocate_tasks
import random
import numpy as np
from model_server import model_predict

def compute_path_length(path):
    if not path or len(path) < 2:
        return 0.0
    total_distance = 0.0
    for i in range(1, len(path)):
        total_distance += np.linalg.norm(np.array(path[i]) - np.array(path[i - 1]))
    return total_distance

def damage_simulate(model, init_res,finsh_rate,damage_rate):
        
        agent_names = list(init_res.keys())
        total_agents = len(agent_names)
        # 每个无人机删除前 40% 的点（视为已完成）
        unvisited_points = []

        # 随机选择损毁的智能体
        damage_ratio = damage_rate
        num_damaged = max(1, int(total_agents * damage_ratio))
        damaged_agents_indices = random.sample(range(total_agents), num_damaged)
        damaged_name = [agent_names[i] for i in damaged_agents_indices]
        for agent_name, path in init_res.items():
            path_len = len(path)
            if path_len <= 1:
                continue

            visited_count = int(path_len * finsh_rate)  # 前40%
            remaining_path = path[visited_count:]
            # 收集该智能体的所有剩余点（未访问点）
            init_res[agent_name] = path[visited_count:]
            if agent_name in damaged_name:
                unvisited_points.extend(remaining_path)
        # 去重（防止同一个点被多个智能体分配）
        unique_unvisited_points = [list(x) for x in set(tuple(point) for point in unvisited_points)]


        # 调用重新分配函数
        updated_agents_positions, elapsed_time = reallocate_tasks(
            init_res, unique_unvisited_points, damaged_agents_indices, model=model
        )
        longest_path_length = 0
        for agent_name, path in updated_agents_positions.items():
            updated_agents_positions[agent_name] = model_predict(model, path, path[0])
            longest_path_length = max(longest_path_length, compute_path_length(path))
        return updated_agents_positions, longest_path_length, elapsed_time
def task_reorganization(model, task_pos, drone_pos, mode_type):
    res_json = {}

    # 初次分配
    init_res = task_allocation_kdtree(model, task_pos,drone_pos)
    res_json['初次分配的结果'] = init_res

    if mode_type == 'static':  # 静态损毁

        updated_agents_positions, longest_path_length, elapsed_time = damage_simulate(model,init_res, 0.6,0.6)
        res_json['静态损毁后重新分配的结果'] = updated_agents_positions
        res_json['静态损毁后重新分配的最大长度'] = longest_path_length
        res_json['静态损毁后重新分配的时间'] = elapsed_time
    elif mode_type == 'dynamic':  # 动态损毁
        updated_agents_positions, longest_path_length, elapsed_time = damage_simulate(model,init_res, 0.2,0.05)
        res_json['第一次动态损毁后重新分配的结果'] = updated_agents_positions
        res_json['第一次动态损毁后重新分配的最大长度'] = longest_path_length
        res_json['第一次动态损毁后重新分配的时间'] = elapsed_time
        updated_agents_positions, longest_path_length, elapsed_time = damage_simulate(model,updated_agents_positions, 0.2,0.1)
        res_json['第二次动态损毁后重新分配的结果'] = updated_agents_positions
        res_json['第二次动态损毁后重新分配的最大长度'] = longest_path_length
        res_json['第二次动态损毁后重新分配的时间'] = elapsed_time
        updated_agents_positions, longest_path_length, elapsed_time = damage_simulate(model,updated_agents_positions, 0.2,0.1)
        res_json['第三次动态损毁后重新分配的结果'] = updated_agents_positions
        res_json['第三次动态损毁后重新分配的最大长度'] = longest_path_length
        res_json['第三次动态损毁后重新分配的时间'] = elapsed_time
    return res_json